#include "Audio_RT_FFT_KF.h"

FFTProcessor::FFTProcessor() {
    initialize_kissfft();
}

FFTProcessor::~FFTProcessor() {
    cleanup_kissfft();
}

void FFTProcessor::initialize_kissfft() 
{
    // 输入缓冲区（实数）
    fft_in_ = new kiss_fft_scalar[FRAME_SIZE];

    // 频域输出（复数：帧大小一半 + 1)
    fft_out_ = new kiss_fft_cpx[FFT_BINS];

    // IFFT 输出缓冲（实数）
    ifft_out_ = new kiss_fft_scalar[FRAME_SIZE];

    // 创建 KissFFT 的实数 FFT 配置
    cfg_forward_  = kiss_fftr_alloc(FRAME_SIZE, 0, nullptr, nullptr);
    cfg_inverse_  = kiss_fftr_alloc(FRAME_SIZE, 1, nullptr, nullptr);
}

void FFTProcessor::cleanup_kissfft() 
{
    if (cfg_forward_) free(cfg_forward_);
    if (cfg_inverse_) free(cfg_inverse_);

    delete[] fft_in_;
    delete[] fft_out_;
    delete[] ifft_out_;
}

void FFTProcessor::compute_spectrum(
    const std::vector<std::vector<__fp16>>& time_data,
    std::vector<std::vector<__fp16>>& mag,
    std::vector<std::vector<__fp16>>& phase)
{
    mag.resize(NUM_FRAMES, std::vector<__fp16>(FFT_BINS));
    phase.resize(NUM_FRAMES, std::vector<__fp16>(FFT_BINS));

    for (size_t i = 0; i < NUM_FRAMES; ++i) {

        // 1. 将 fp16 转换为 float 写入 KissFFT 输入
        for (size_t j = 0; j < FRAME_SIZE; ++j) {
            fft_in_[j] = static_cast<float>(time_data[i][j]);
        }

        // 2. 执行 KissFFT R2C
        kiss_fftr(cfg_forward_, fft_in_, fft_out_);

        // 3. 计算幅度 + 相位
        for (size_t k = 0; k < FFT_BINS; ++k) {
            float real = fft_out_[k].r;
            float imag = fft_out_[k].i;

            float magnitude = std::sqrt(real * real + imag * imag);
            mag[i][k] = static_cast<__fp16>(magnitude);

            // 相位归一化 [-π, π] → [-1, 1]
            phase[i][k] = static_cast<__fp16>(std::atan2(imag, real) / M_PI);
        }
    }
}


// 进行单独一帧的FFT，以适应新数据帧的移入
void FFTProcessor::compute_spectrum_One(
    const std::vector<std::vector<__fp16>>& time_data,
    std::vector<std::vector<__fp16>>& mag,
    std::vector<std::vector<__fp16>>& phase)
{
    // 1. 拷贝最后一帧的数据到 FFT 输入
    for (size_t j = 0; j < FRAME_SIZE; ++j) {
        fft_in_[j] = static_cast<float>(time_data[NUM_FRAMES - 1][j]);
    }

    // 2. 实数 FFT
    kiss_fftr(cfg_forward_, fft_in_, fft_out_);

    // 3. 计算幅度和相位
    for (size_t k = 0; k < FFT_BINS; ++k) {
        float real = fft_out_[k].r;
        float imag = fft_out_[k].i;

        float magnitude = std::sqrt(real * real + imag * imag);
        mag[NUM_FRAMES - 1][k] = static_cast<__fp16>(magnitude);

        // 相位归一化 [-π,π] → [-1,1]
        phase[NUM_FRAMES - 1][k] = static_cast<__fp16>(std::atan2(imag, real) / M_PI);
    }
}

// 只对mag和phase的最后两行进行IFFT
void FFTProcessor::reconstruct_signal_Only(
    const std::vector<std::vector<__fp16>>& mag,
    const std::vector<std::vector<__fp16>>& phase,
    std::vector<std::vector<__fp16>>& time_data)
{
    // 固定最后2帧
    time_data.resize(2, std::vector<__fp16>(FRAME_SIZE));

    size_t num_frames = mag.size(); // 动态获取总帧数

    for (size_t i = num_frames - 2; i < num_frames; ++i) {

        // 1. 重建复数谱
        for (size_t k = 0; k < FFT_BINS; ++k) {
            float magnitude = static_cast<float>(mag[i][k]);
            float angle     = static_cast<float>(phase[i][k]) * M_PI;

            fft_out_[k].r = magnitude * std::cos(angle);
            fft_out_[k].i = magnitude * std::sin(angle);
        }

        // 2. 执行 IFFT（KissFFT）
        kiss_fftri(cfg_inverse_, fft_out_, ifft_out_);

        // 3. 缩放并转换回 fp16
        const float scale = 1.0f / FRAME_SIZE;
        for (size_t j = 0; j < FRAME_SIZE; ++j) {
            time_data[i - num_frames + 2][j] = static_cast<__fp16>(ifft_out_[j] * scale);
        }
    }
}


void FFTProcessor::reconstruct_signal(
    const std::vector<std::vector<__fp16>>& mag,
    const std::vector<std::vector<__fp16>>& phase,
    std::vector<std::vector<__fp16>>& time_data)
{
    time_data.resize(NUM_FRAMES, std::vector<__fp16>(FRAME_SIZE));

    for (size_t i = 0; i < NUM_FRAMES; ++i) {

        // 1. 重建复数谱
        for (size_t k = 0; k < FFT_BINS; ++k) {
            float magnitude = static_cast<float>(mag[i][k]);
            float angle     = static_cast<float>(phase[i][k]) * M_PI;

            fft_out_[k].r = magnitude * std::cos(angle);
            fft_out_[k].i = magnitude * std::sin(angle);
        }

        // 2. 执行逆 FFT
        kiss_fftri(cfg_inverse_, fft_out_, ifft_out_);

        // 3. 缩放并转换回 fp16
        const float scale = 1.0f / FRAME_SIZE;
        for (size_t j = 0; j < FRAME_SIZE; ++j) {
            time_data[i][j] = static_cast<__fp16>(ifft_out_[j] * scale);
        }
    }
}